import base64
import sys
from Crypto.Cipher import AES


# 加密
class Encrypt():
    def __init__(self, key):  # self为EncryptData类的对象
        self.key = key.encode("utf-8")  # 初始化秘钥
        self.length = AES.block_size  # 初始化数据块大小 16
        self.aes = AES.new(self.key, AES.MODE_ECB)  # 初始化AES，ECB模式
        # 删除加密之后填充的字符 用于解密
        self.unpad = lambda date: date[0:-ord(date[-1])]

    def pad(self, text):
        # 使被加密数据的字节码长度是block_size的整数倍
        count = len(text.encode('utf-8'))
        add = self.length - (count % self.length)
        entext = text + (chr(add) * add)
        return entext

    def encrypt(self, encrData):  # 加密函数
        res = self.aes.encrypt(self.pad(encrData).encode("utf8"))
        msg = str(base64.b64encode(res), encoding="utf8")
        return msg


if __name__ == '__main__':
    # 命令行：python test.py parameter1 parameter2
    key = sys.argv[1]  # key 自己输入的秘钥  获取命令行的第一个参数
    data = sys.argv[2]  # 要加密的数据   获取命令行的第二个参数
    eg = Encrypt(key)  # 这里秘钥长度必须是16
    # 加密
    res = eg.encrypt(str(data))
    print("加密后：" + res, end='')
